package com.bxmdm.ragdemo.reader;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.core.io.Resource;
import org.springframework.util.StreamUtils;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * @Description: 英语问答数据读取器，简化版本 - 只进行问题向量化
 */
public class EnglishQAReader implements DocumentReader {

    private final Resource resource;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final EmbeddingClient embeddingClient;

    // 文档缓存
    private List<Document> documentCache = null;

    public EnglishQAReader(Resource resource, EmbeddingClient embeddingClient) {
        this.resource = resource;
        this.embeddingClient = embeddingClient;
    }

    @Override
    public List<Document> get() {
        // 检查缓存
        if (documentCache != null) {
            return documentCache;
        }

        try {
            List<Document> documents = new ArrayList<>();
            String content = StreamUtils.copyToString(this.resource.getInputStream(), StandardCharsets.UTF_8);

            // 从JSON中解析问答列表
            List<Map<String, Object>> qaItems = objectMapper.readValue(
                    content,
                    new TypeReference<List<Map<String, Object>>>() {
                    });

            for (Map<String, Object> qaItem : qaItems) {
                Document doc = processQAPair(qaItem);
                if (doc != null) {
                    documents.add(doc);
                }
            }

            // 缓存结果
            documentCache = documents;
            return documents;
        } catch (IOException e) {
            throw new RuntimeException("读取问答数据文件失败: " + e.getMessage(), e);
        }
    }

    @SuppressWarnings("unchecked")
    private Document processQAPair(Map<String, Object> qaItem) {
        // 提取基本信息
        String postId = qaItem.getOrDefault("post_id", "").toString();
        String question = qaItem.getOrDefault("question", "").toString();
        String answer = qaItem.getOrDefault("answer", "").toString();
        Integer score = qaItem.containsKey("score") ? Integer.valueOf(qaItem.get("score").toString()) : 0;
        List<String> tags = qaItem.containsKey("tags") ? (List<String>) qaItem.get("tags") : new ArrayList<>();
        Integer answerCount = qaItem.containsKey("answer_count")
                ? Integer.valueOf(qaItem.get("answer_count").toString())
                : 0;
        List<String> allAnswers = qaItem.containsKey("all_answers") ? (List<String>) qaItem.get("all_answers")
                : new ArrayList<>();

        // 如果问题为空，跳过
        if (question.isEmpty()) {
            return null;
        }

        // 提取问题的第一行作为主要问题
        String[] lines = question.split("\\r?\\n");
        String mainQuestion = lines[0].trim();

        // 使用加权向量计算策略 - 对完整问题内容进行处理
        List<Double> embedding = null;
        try {
            embedding = calculateWeightedEmbedding(question); // 对完整问题进行加权嵌入
        } catch (Exception e) {
            System.err.println("计算问题嵌入向量失败: " + e.getMessage());
            return null;
        }

        // 同时保存增强的文本内容用于调试
        String enhancedQuestionContent = enhanceQuestionSemantics(question);

        // 构建元数据
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("post_id", postId);
        metadata.put("score", score);
        metadata.put("question", question); // 保存完整问题
        metadata.put("main_question", mainQuestion); // 保存主要问题
        metadata.put("enhanced_content", enhancedQuestionContent); // 保存增强内容
        metadata.put("answer", answer);
        metadata.put("tags", tags);
        metadata.put("answer_count", answerCount);
        metadata.put("all_answers", allAnswers);
        metadata.put("embedding", embedding);

        // 创建Document对象 - 使用完整问题内容作为文档内容
        String id = UUID.randomUUID().toString();
        Document document = new Document(id, question, metadata);
        return document;
    }

    /**
     * 计算加权向量嵌入 - 改进版本
     * 使用混合策略：结合整体嵌入和加权句子嵌入
     */
    private List<Double> calculateWeightedEmbedding(String question) {
        // 1. 清理和标准化问题文本
        String cleanQuestion = question.replaceAll("\\s+", " ").trim();

        // 2. 首先计算整体问题的嵌入作为基础
        List<Double> globalEmbedding = embeddingClient.embed(cleanQuestion);

        // 3. 分句处理 - 如果问题很短或只有一句，直接返回全局嵌入
        String[] sentences = cleanQuestion.split("[.!?\\n]+");
        if (sentences.length <= 1 || cleanQuestion.length() < 50) {
            return globalEmbedding;
        }

        List<List<Double>> sentenceEmbeddings = new ArrayList<>();
        List<Double> weights = new ArrayList<>();

        // 4. 为每个句子计算向量和权重
        for (int i = 0; i < sentences.length; i++) {
            String sentence = sentences[i].trim();
            if (sentence.isEmpty() || sentence.length() < 10) {
                continue;
            }

            // 计算句子的向量嵌入
            List<Double> sentenceEmbedding = embeddingClient.embed(sentence);
            sentenceEmbeddings.add(sentenceEmbedding);

            // 计算权重
            double weight = calculateSentenceWeight(sentence, i, sentences.length);
            weights.add(weight);
        }

        // 5. 如果没有有效句子或只有一个句子，返回全局嵌入
        if (sentenceEmbeddings.isEmpty() || sentenceEmbeddings.size() == 1) {
            return globalEmbedding;
        }

        // 6. 计算加权平均向量
        List<Double> weightedEmbedding = calculateWeightedAverage(sentenceEmbeddings, weights);

        // 7. 混合策略：将全局嵌入和加权嵌入按比例混合
        // 全局嵌入权重为0.4，加权嵌入权重为0.6
        return combineEmbeddings(globalEmbedding, weightedEmbedding, 0.4, 0.6);
    }

    /**
     * 计算句子权重 - 改进版本
     * 
     * @param sentence       句子内容
     * @param index          句子在问题中的位置索引
     * @param totalSentences 总句子数
     * @return 句子权重
     */
    private double calculateSentenceWeight(String sentence, int index, int totalSentences) {
        double weight = 1.0;

        // 1. 位置权重策略调整 - 降低第一句话的过高权重
        if (index == 0) {
            weight = 1.8; // 第一句话权重降低为1.8倍
        } else {
            // 其他句子权重更平缓的递减
            weight = 1.0 - (index * 0.05); // 每后一句减少0.05权重
            weight = Math.max(weight, 0.5); // 最低权重为0.5
        }

        // 2. 关键语义句子权重调整
        if (isKeySemanticSentence(sentence)) {
            weight *= 1.3; // 降低语义加权倍数
        }

        // 3. 问号结尾的句子 - 这是最重要的特征
        if (sentence.trim().endsWith("?")) {
            weight *= 1.5; // 增加问句权重
        }

        // 4. 包含关键比较词的句子增加权重
        String lower = sentence.toLowerCase();
        if (lower.contains("difference") || lower.contains("vs") || lower.contains("versus")
                || lower.contains("between") || lower.contains("or")) {
            weight *= 1.4;
        }

        // 5. 包含引号的句子增加权重（通常包含具体例子）
        if (sentence.contains("\"") && sentence.contains("\"")) {
            weight *= 1.2;
        }

        // 6. 包含模态动词的句子增加权重
        String[] modalVerbs = { "can", "could", "will", "would", "shall", "should", "may", "might", "must" };
        for (String modal : modalVerbs) {
            if (lower.contains(" " + modal + " ") || lower.startsWith(modal + " ")) {
                weight *= 1.2;
                break;
            }
        }

        return weight;
    }

    /**
     * 计算加权平均向量
     * 
     * @param embeddings 向量列表
     * @param weights    权重列表
     * @return 加权平均向量
     */
    private List<Double> calculateWeightedAverage(List<List<Double>> embeddings, List<Double> weights) {
        if (embeddings.isEmpty() || weights.isEmpty() || embeddings.size() != weights.size()) {
            throw new IllegalArgumentException("向量列表和权重列表不匹配");
        }

        // 获取向量维度
        int dimension = embeddings.get(0).size();

        // 初始化结果向量
        List<Double> result = new ArrayList<>(dimension);
        for (int i = 0; i < dimension; i++) {
            result.add(0.0);
        }

        // 计算权重总和
        double totalWeight = weights.stream().mapToDouble(Double::doubleValue).sum();

        // 计算加权平均
        for (int i = 0; i < embeddings.size(); i++) {
            List<Double> embedding = embeddings.get(i);
            double weight = weights.get(i);
            double normalizedWeight = weight / totalWeight;

            for (int j = 0; j < dimension; j++) {
                double currentValue = result.get(j);
                double weightedValue = embedding.get(j) * normalizedWeight;
                result.set(j, currentValue + weightedValue);
            }
        }

        return result;
    }

    /**
     * 增强问题语义提取 - 改进版本
     * 对于英语学习问答，更有针对性地提取和增强关键概念
     */
    private String enhanceQuestionSemantics(String question) {
        StringBuilder enhanced = new StringBuilder();

        // 1. 清理和标准化问题文本
        String cleanQuestion = question.replaceAll("\\s+", " ").trim();

        // 2. 保持原始问题结构
        enhanced.append(cleanQuestion);

        // 3. 提取并强化关键语法概念
        String grammarConcepts = extractGrammarConcepts(cleanQuestion);
        if (!grammarConcepts.isEmpty()) {
            enhanced.append(" [Grammar: ").append(grammarConcepts).append("]");
        }

        // 4. 提取比较对象
        String comparisonTerms = extractComparisonTerms(cleanQuestion);
        if (!comparisonTerms.isEmpty()) {
            enhanced.append(" [Compare: ").append(comparisonTerms).append("]");
        }

        // 5. 提取使用场景关键词
        String contextTerms = extractContextTerms(cleanQuestion);
        if (!contextTerms.isEmpty()) {
            enhanced.append(" [Context: ").append(contextTerms).append("]");
        }

        return enhanced.toString().trim();
    }

    /**
     * 提取语法概念关键词
     */
    private String extractGrammarConcepts(String question) {
        StringBuilder concepts = new StringBuilder();
        String lower = question.toLowerCase();

        // 语法类别
        String[] grammarCategories = {
                "modal", "verb", "tense", "article", "preposition", "pronoun",
                "adjective", "adverb", "conjunction", "singular", "plural",
                "past", "present", "future", "perfect", "progressive", "conditional"
        };

        for (String category : grammarCategories) {
            if (lower.contains(category)) {
                concepts.append(category).append(" ");
            }
        }

        // 特定的语法结构
        if (lower.contains("passive"))
            concepts.append("passive_voice ");
        if (lower.contains("active"))
            concepts.append("active_voice ");
        if (lower.contains("subjunctive"))
            concepts.append("subjunctive_mood ");
        if (lower.contains("imperative"))
            concepts.append("imperative_mood ");

        return concepts.toString().trim();
    }

    /**
     * 提取比较术语
     */
    private String extractComparisonTerms(String question) {
        StringBuilder terms = new StringBuilder();
        String lower = question.toLowerCase();

        // 提取引号中的词汇
        java.util.regex.Pattern quotePattern = java.util.regex.Pattern.compile("\"([^\"]+)\"");
        java.util.regex.Matcher quoteMatcher = quotePattern.matcher(question);
        while (quoteMatcher.find()) {
            terms.append(quoteMatcher.group(1)).append(" ");
        }

        // 提取vs, versus, or连接的词汇
        java.util.regex.Pattern vsPattern = java.util.regex.Pattern.compile("(\\w+)\\s+(?:vs|versus|or)\\s+(\\w+)");
        java.util.regex.Matcher vsMatcher = vsPattern.matcher(lower);
        while (vsMatcher.find()) {
            terms.append(vsMatcher.group(1)).append(" ").append(vsMatcher.group(2)).append(" ");
        }

        // 提取between...and...结构
        java.util.regex.Pattern betweenPattern = java.util.regex.Pattern.compile("between\\s+(\\w+)\\s+and\\s+(\\w+)");
        java.util.regex.Matcher betweenMatcher = betweenPattern.matcher(lower);
        while (betweenMatcher.find()) {
            terms.append(betweenMatcher.group(1)).append(" ").append(betweenMatcher.group(2)).append(" ");
        }

        return terms.toString().trim();
    }

    /**
     * 提取使用场景关键词
     */
    private String extractContextTerms(String question) {
        StringBuilder terms = new StringBuilder();
        String lower = question.toLowerCase();

        String[] contextKeywords = {
                "formal", "informal", "spoken", "written", "british", "american",
                "appropriate", "correct", "proper", "acceptable", "polite",
                "situation", "context", "time", "place", "way"
        };

        for (String keyword : contextKeywords) {
            if (lower.contains(keyword)) {
                terms.append(keyword).append(" ");
            }
        }

        return terms.toString().trim();
    }

    /**
     * 判断是否为关键语义句子 - 改进版本
     * 更精确地识别英语语法、用法问题的关键句子
     */
    private boolean isKeySemanticSentence(String sentence) {
        String lower = sentence.toLowerCase();

        // 1. 疑问词开头的句子（核心问题）
        String[] questionWords = { "when", "how", "what", "which", "why", "where", "who" };
        for (String qw : questionWords) {
            if (lower.startsWith(qw + " ")) {
                return true;
            }
        }

        // 2. 核心语法问题特征词汇
        String[] coreGrammarKeywords = {
                "difference", "should", "correct", "appropriate", "proper", "right", "better",
                "use", "mean", "versus", "vs", "between", "or"
        };
        for (String keyword : coreGrammarKeywords) {
            if (lower.contains(keyword)) {
                return true;
            }
        }

        // 3. 模态动词相关问题
        String[] modalVerbs = { "can", "could", "will", "would", "shall", "should", "may", "might", "must" };
        for (String modal : modalVerbs) {
            if (lower.contains(modal)) {
                return true;
            }
        }

        // 4. 包含引号的句子（通常是具体的语法例子）
        if (sentence.contains("\"") && sentence.contains("\"")) {
            return true;
        }

        // 5. 时态和语法结构相关
        String[] grammarStructures = {
                "tense", "passive", "active", "past", "present", "future",
                "singular", "plural", "article", "preposition"
        };
        for (String structure : grammarStructures) {
            if (lower.contains(structure)) {
                return true;
            }
        }

        // 6. 使用场景相关
        String[] contextWords = {
                "formal", "informal", "british", "american", "spoken", "written",
                "situation", "context", "appropriate", "acceptable"
        };
        for (String context : contextWords) {
            if (lower.contains(context)) {
                return true;
            }
        }

        return false;
    }

    /**
     * 组合两个嵌入向量
     * 
     * @param embedding1 第一个嵌入向量
     * @param embedding2 第二个嵌入向量
     * @param weight1    第一个向量的权重
     * @param weight2    第二个向量的权重
     * @return 组合后的嵌入向量
     */
    private List<Double> combineEmbeddings(List<Double> embedding1, List<Double> embedding2,
            double weight1, double weight2) {
        if (embedding1.size() != embedding2.size()) {
            throw new IllegalArgumentException("两个嵌入向量的维度必须相同");
        }

        List<Double> result = new ArrayList<>();
        for (int i = 0; i < embedding1.size(); i++) {
            double combinedValue = embedding1.get(i) * weight1 + embedding2.get(i) * weight2;
            result.add(combinedValue);
        }
        return result;
    }
}